package cc.shacocloud.luanniao.core.service.ai

import org.springframework.ai.chat.model.ChatModel
import org.springframework.ai.chat.model.ChatResponse
import org.springframework.ai.chat.prompt.Prompt
import org.springframework.ai.document.Document
import org.springframework.ai.embedding.EmbeddingModel
import org.springframework.ai.embedding.EmbeddingRequest
import org.springframework.ai.embedding.EmbeddingResponse
import reactor.core.publisher.Flux

/**
 * 抽象的 ai 模型服务
 * @author 思追(shaco)
 */
abstract class AbstractAiModel(
    val chatModel: ChatModel,
    val embeddingModel: EmbeddingModel,
) : AiModel, AutoCloseable {

    override fun chatCall(prompt: Prompt): ChatResponse {
        return chatModel.call(prompt)
    }

    override fun chatStream(prompt: Prompt): Flux<ChatResponse> {
        return chatModel.stream(prompt)
    }

    override fun embedCall(request: EmbeddingRequest): EmbeddingResponse {
        return embeddingModel.call(request)
    }

    override fun embed(document: Document): FloatArray {
        return embeddingModel.embed(document)
    }

    override fun close() {
        if (chatModel is AutoCloseable) {
            chatModel.close()
        }

        if (embeddingModel is AutoCloseable) {
            embeddingModel.close()
        }
    }
}